package cn.hc.tool.concurrent.util;

import cn.hc.tool.common.util.CollectionUtil;
import cn.hc.tool.common.util.MapUtil;
import cn.hc.tool.trace.common.HcTraceConst;
import lombok.Getter;
import lombok.extern.slf4j.Slf4j;
import org.slf4j.MDC;
import org.springframework.beans.factory.InitializingBean;
import org.springframework.beans.factory.annotation.Autowired;

import java.util.*;
import java.util.concurrent.Callable;
import java.util.concurrent.ExecutorService;
import java.util.concurrent.Future;
import java.util.function.BiFunction;
import java.util.function.Function;

/**
 * @author huangchao E-mail:fengquan8866@163.com
 * @version 创建时间：2024/10/17 19:34
 */
@Slf4j
public class HcParallelUtil implements InitializingBean {

    /**
     * 子线程需要同步的ThreadLocal对象
     */
    private final Set<ThreadLocal<?>> threadLocals = new HashSet<>();

    /**
     * 子线程需要同步的MDC数据
     */
    private final Set<String> mdcKeys = new HashSet<>();

    /**
     * 线程池
     */
    @Getter
    private final ExecutorService executor;

    /**
     * 让 HcTraceConst 先初始化
     */
    @Autowired(required = false)
    private HcTraceConst traceConst;
//
//    /**
//     * 默认线程池构造器
//     */
//    public ToolParallelUtil() {
//        executor = new ThreadPoolExecutor(100, 100,
//                5, TimeUnit.SECONDS, new LinkedBlockingQueue<>(2000), r -> {
//            Thread thread = new Thread(r);
//            thread.setName("ParallelCompUtil-" + thread.hashCode());
//            return thread;
//        });
//    }

    /**
     * 自定义线程池构造器
     *
     * @param executor 自定义线程池
     */
    public HcParallelUtil(ExecutorService executor) {
        this.executor = executor;
    }

    /**
     * {@link #runAsync(ExecutorService, Runnable...)}
     */
    public final void runAsync(Runnable... runnable) {
        this.runAsync(null, runnable);
    }

    /**
     * 并行执行异步方法
     *
     * @param executor 自定义线程池
     * @param runnable 异步方法列表
     */
    private void runAsync(ExecutorService executor, Runnable... runnable) {
        long tid = Thread.currentThread().getId();
        Map<Integer, Object> threadLocalMap = getThreadLocalMap();
        Map<String, String> mdcMap = getMdcMap();
        for (Runnable a : runnable) {
            Runnable runnable1 = this.generateContextRunnable(a, tid, threadLocalMap, mdcMap);
            this.executeAll(executor, runnable1);
        }
    }

    /**
     * {@link #runAsync(ExecutorService, Callable[])}
     */
    @SafeVarargs
    public final <V> List<Future<V>> runAsync(Callable<V>... callables) {
        return runAsync(null, callables);
    }

    /**
     * 并行执行异步方法，返回Future集合
     *
     * @param executor 自定义线程池
     * @param callables 异步方法列表
     * @param <V> 异步方法返回值类型
     * @return 异步方法执行结果列表
     */
    @SafeVarargs
    public final <V> List<Future<V>> runAsync(ExecutorService executor, Callable<V>... callables) {
        long tid = Thread.currentThread().getId();
        Map<Integer, Object> threadLocalMap = getThreadLocalMap();
        Map<String, String> mdcMap = getMdcMap();
        List<Callable<V>> ls = new ArrayList<>();
        for (Callable<V> a : callables) {
            ls.add(generateContextCallable(a, tid, threadLocalMap, mdcMap));
        }
        return this.submitAll(executor, ls);
    }

    /**
     * 生成包含上下文的Callable
     *
     * @param callable 原始callable
     * @param <V>      callable泛型
     */
    public <V> Callable<V> generateContextCallable(Callable<V> callable) {
        long tid = Thread.currentThread().getId();
        Map<Integer, Object> threadLocalMap = getThreadLocalMap();
        Map<String, String> mdcMap = getMdcMap();
        return this.generateContextCallable(callable, tid, threadLocalMap, mdcMap);
    }

    /**
     * future get
     *
     * @param futures future列表
     * @return 任务执行结果列表
     */
    public <V> List<V> get(List<Future<V>> futures) {
        List<V> list = new ArrayList<>();
        for (Future<V> f : futures) {
            try {
                list.add(f.get());
            } catch (Exception e) {
                log.error("error in get futures", e);
                list.add(null);
            }
        }
        return list;
    }

    /**
     * {@link #complete(ExecutorService, Callable[])}
     */
    @SafeVarargs
    public final <V> List<V> complete(Callable<V>... callables) {
        return complete(null, callables);
    }

    /**
     * 完成处理多个异步方法
     *
     * @param executor 自定义线程池
     * @param callables 任务列表
     * @return 任务执行结果列表
     */
    @SafeVarargs
    public final <V> List<V> complete(ExecutorService executor, Callable<V>... callables) {
        List<Future<V>> list = runAsync(executor, callables);
        return this.get(list);
    }

    /**
     * 并发执行
     *
     * @param params 参数
     * @param call   回调函数
     * @param <P>    入参类型
     * @param <V>    出参类型
     * @return 执行结果集
     */
    public <P, V> List<V> allOf(Collection<P> params, Function<P, V> call) {
        return allOf(null, params, call);
    }

    public <P, V> List<V> allOf(ExecutorService executor, Collection<P> params, Function<P, V> call) {
        int size = params.size();
        // 单条数据，直接执行
        if (size == 1) {
            P p = null;
            for (P p1 : params) {
                p = p1;
            }
            List<V> result = new ArrayList<>();
            result.add(call.apply(p));
            return result;
        }
        // 多条数据，走线程池执行
        else {
            Callable<V>[] callables = generateCallable(params, call);
            return complete(executor, callables);
        }
    }

    /**
     * {@link #allOfRun(ExecutorService, Collection, Function)}
     */
    public <P, V> List<Future<V>> allOfRun(Collection<P> params, Function<P, V> call) {
        return allOfRun(null, params, call);
    }

    /**
     * 并发执行
     *
     * @param executor 自定义线程池
     * @param params   参数列表
     * @param call     执行体
     * @param <P>      入参类型
     * @param <V>      出参类型
     * @return Future对象列表
     */
    public <P, V> List<Future<V>> allOfRun(ExecutorService executor, Collection<P> params, Function<P, V> call) {
        Callable<V>[] callables = this.generateCallable(params, call);
        return this.runAsync(executor, callables);
    }

    /**
     * {@link #allOfRun(ExecutorService, Collection, Function, Map)}
     */
    public <P, V> List<Future<V>> allOfRun(Collection<P> params, Function<P, V> call, Map<P, V> map) {
        return allOfRun(null, params, call, map);
    }

    /**
     * 并发执行
     *
     * @param executor     执行器
     * @param params       参数列表
     * @param call         执行体
     * @param map          构造结果的map
     * @param <P>          入参类型
     * @param <V>          出参类型
     * @return Future对象列表
     */
    public <P, V> List<Future<V>> allOfRun(ExecutorService executor, Collection<P> params, Function<P, V> call,
                                           Map<P, V> map) {
        return allOfRun(executor, params, call, (p, v) -> p, map);
    }

    /**
     * {@link #allOfRun(ExecutorService, Collection, Function, BiFunction, Map)}
     */
    public <P, V, K> List<Future<V>> allOfRun(Collection<P> params, Function<P, V> call,
                                              BiFunction<P, V, K> keyGenerator, Map<K, V> map) {
        return allOfRun(null, params, call, keyGenerator, map);
    }

    /**
     * 并发执行
     *
     * @param executor     执行器
     * @param params       参数列表
     * @param call         执行体
     * @param keyGenerator map的key生成器
     * @param map          构造结果的map
     * @param <P>          入参类型
     * @param <V>          出参类型
     * @param <K>          map的key类型
     * @return Future对象列表
     */
    public <P, V, K> List<Future<V>> allOfRun(ExecutorService executor, Collection<P> params, Function<P, V> call,
                                              BiFunction<P, V, K> keyGenerator, Map<K, V> map) {
        Callable<V>[] callables = generate2MapCallable(params, call, keyGenerator, map);
        return runAsync(executor, callables);
    }

    /**
     * 添加ThreadLocal对象
     */
    public void addThreadLocals(ThreadLocal... threadLocals) {
        log.info("添加ThreadLocal：{}", Arrays.toString(threadLocals));
        for (ThreadLocal<?> t : threadLocals) {
            if (t != null) {
                this.threadLocals.add(t);
            }
        }
    }

    /**
     * 添加MDC数据
     */
    public void addMdcKey(String... mdcKeys) {
        log.info("添加mdcKey：{}", Arrays.toString(mdcKeys));
        for (String t : mdcKeys) {
            if (t != null) {
                this.mdcKeys.add(t);
            }
        }
    }

    /**
     * 任务转换生成
     *
     * @param params 参数列表
     * @param call   任务执行体
     */
    private <P, V> Callable<V>[] generateCallable(Collection<P> params, Function<P, V> call) {
        Callable<V>[] callables = new Callable[params.size()];
        int i = 0;
        for (P p : params) {
            callables[i++] = () -> call.apply(p);
        }
        return callables;
    }

    /**
     * 生成到Map的任务 转换生成
     *
     * @param params       参数列表
     * @param call         任务执行体
     * @param keyGenerator map的key生成器
     * @param resMap       返回的map
     */
    private <P, V, K> Callable<V>[] generate2MapCallable(Collection<P> params, Function<P, V> call,
                                                         BiFunction<P, V, K> keyGenerator, Map<K, V> resMap) {
        Callable<V>[] callables = new Callable[params.size()];
        int i = 0;
        for (P p : params) {
            callables[i++] = () -> {
                V res = call.apply(p);
                resMap.put(keyGenerator.apply(p, res), res);
                return res;
            };
        }
        return callables;
    }

    /**
     * 任务放入线程池
     *
     * @param executor 执行器
     * @param ls       任务列表
     */
    private <C extends Runnable> void executeAll(ExecutorService executor, C... ls) {
        if (executor == null) {
            executor = this.executor;
        }
        for (C c : ls) {
            executor.execute(c);
        }
    }

    /**
     * 任务放入线程池
     *
     * @param executor 执行器
     * @param ls       任务列表
     * @param <V>      返回数据类型
     * @return Future列表
     */
    private <V, C extends Callable<V>> List<Future<V>> submitAll(ExecutorService executor, Collection<C> ls) {
        if (executor == null) {
            executor = this.executor;
        }
        List<Future<V>> result = new ArrayList<>();
        for (Callable<V> c : ls) {
            result.add(executor.submit(c));
        }
        return result;
    }

    private Runnable generateContextRunnable(Runnable runnable, long tid,
                                             Map<Integer, Object> threadLocalMap, Map<String, String> mdcMap) {
        return () -> {
            try {
                if (tid != Thread.currentThread().getId()) {
                    initThreadLocal(threadLocalMap);
                    initMdc(mdcMap);
                }
                runnable.run();
            } finally {
                if (tid != Thread.currentThread().getId()) {
                    clearThreadLocal(threadLocalMap);
                    clearMdc(mdcMap);
                }
            }
        };
    }

    private <V> Callable<V> generateContextCallable(Callable<V> callable, long tid,
                                                    Map<Integer, Object> threadLocalMap, Map<String, String> mdcMap) {
        return () -> {
            try {
                if (tid != Thread.currentThread().getId()) {
                    initThreadLocal(threadLocalMap);
                    initMdc(mdcMap);
                }
                return callable.call();
            } finally {
                if (tid != Thread.currentThread().getId()) {
                    clearThreadLocal(threadLocalMap);
                    clearMdc(mdcMap);
                }
            }
        };
    }

    /**
     * 清理子线程MDC
     *
     * @param mdcMap 父线程MDC数据
     */
    private void clearMdc(Map<String, String> mdcMap) {
        if (MapUtil.isEmpty(mdcMap)) {
            return;
        }
        for (String k : mdcMap.keySet()) {
            MDC.remove(k);
        }
    }

    /**
     * 初始化子线程MDC数据
     *
     * @param mdcMap 父线程MDC数据
     */
    private void initMdc(Map<String, String> mdcMap) {
        if (MapUtil.isEmpty(mdcMap)) {
            return;
        }
        for (Map.Entry<String, String> e : mdcMap.entrySet()) {
            MDC.put(e.getKey(), e.getValue());
        }
    }

    /**
     * 获取MDC数据
     */
    private Map<String, String> getMdcMap() {
        if (CollectionUtil.isEmpty(mdcKeys)) {
            return null;
        }
        Map<String, String> map = new HashMap<>();
        for (String k : mdcKeys) {
            map.put(k, MDC.get(k));
        }
        return map;
    }

    /**
     * 清理子线程ThreadLocal
     *
     * @param map 父线程数据
     */
    private void clearThreadLocal(Map<Integer, Object> map) {
        if (MapUtil.isEmpty(map)) {
            return;
        }
        for (ThreadLocal<?> t : threadLocals) {
            t.remove();
        }
    }

    /**
     * 初始化子线程ThreadLocal
     *
     * @param map 父线程数据
     */
    private void initThreadLocal(Map<Integer, Object> map) {
        if (MapUtil.isEmpty(map)) {
            return;
        }
        for (ThreadLocal t : threadLocals) {
            t.set(map.get(t.hashCode()));
        }
    }

    /**
     * 获取ThreadLocal数据
     */
    private Map<Integer, Object> getThreadLocalMap() {
        if (CollectionUtil.isEmpty(threadLocals)) {
            return null;
        }
        Map<Integer, Object> map = new HashMap<>();
        for (ThreadLocal<?> t : threadLocals) {
            map.put(t.hashCode(), t.get());
        }
        return map;
    }

    @Override
    public void afterPropertiesSet() throws Exception {
        this.addMdcKey(HcTraceConst.TRACE_ID);
    }
}
